-- file: chap05/Prettify.hs
module Prettify where

import Data.Bits (shiftR, (.&.))
import Data.Char (ord)
import Numeric (showHex)
import Prelude hiding ((<>))

-- ---------- Doc 类型
data Doc
  = Empty
  | Char Char
  | Text String
  | Line -- 换行
  | Concat Doc Doc -- 连接
  | Union Doc Doc -- flatten 形式 和 换行形式
  deriving (Show, Eq)

-- 字符串化 JObject field 的 name
string :: String -> Doc
string "" = Empty
string str = Text (show str)

text :: String -> Doc
text "" = Empty
text s = Text s

double :: Double -> Doc
double d = text (show d)

empty :: Doc
empty = Empty

char :: Char -> Doc
char = Char

line :: Doc
line = Line

(<>) :: Doc -> Doc -> Doc
Empty <> y = y
x <> Empty = x
x <> y = x `Concat` y

-- ------------------ 美观打印
enclose :: Char -> Char -> Doc -> Doc
enclose left right x = char left <> x <> char right

-- 打印单个字符
oneChar :: Char -> Doc
oneChar c = case lookup c simpleEscapes of
  Just r -> text r
  Nothing
    | mustEscape c -> hexEscape c
    | otherwise -> char c
  where
    mustEscape c = c < ' ' || c == '\x7f' || c > '\xff'

simpleEscapes :: [(Char, String)]
simpleEscapes = zipWith ch "\b\n\f\r\t\\\"/" "bnfrt\\\"/"
  where
    ch a b = (a, ['\\', b])

smallHex :: Int -> Doc
smallHex x =
  text "\\u"
    <> text (replicate (4 - length h) '0')
    <> text h
  where
    h = showHex x ""

astral :: Int -> Doc
-- 0xd800 = 1101 1000 0000 0000
-- 0xdc00 = 1101 1100 0000 0000
astral n = smallHex (a + 0xd800) <> smallHex (b + 0xdc00)
  where
    -- 0x3ff 就是 0b11 1111 1111
    a = (n `shiftR` 10) .&. 0x3ff
    b = n .&. 0x3ff

hexEscape :: Char -> Doc
hexEscape c
  | d < 0x10000 = smallHex d
  | otherwise = astral (d - 0x10000)
  where
    d = ord c

-- 美观打印数组和对象就简单多了
series :: Char -> Char -> (a -> Doc) -> [a] -> Doc
series open close f =
  enclose open close
    . fsep
    . punctuate (char ',')
    . map f

-- 加标点符号; 例如加逗号
punctuate :: Doc -> [Doc] -> [Doc]
punctuate p [] = []
punctuate p [d] = [d]
punctuate p (d : ds) = (d <> p) : punctuate p ds

-- 我们的 hcat 和 fsep 函数将 Doc 列表拼接成一个 Doc 值
hcat :: [Doc] -> Doc
hcat = fold (<>)

fold :: (Doc -> Doc -> Doc) -> [Doc] -> Doc
fold f = foldr f empty

-- 它将多个 Doc 值拼接成一个，并且在需要的时候换行
fsep :: [Doc] -> Doc
fsep = fold (</>)

(</>) :: Doc -> Doc -> Doc
x </> y = x <> softline <> y

softline :: Doc
softline = group line

-- file: ch05/Prettify.hs
group :: Doc -> Doc
group x = flatten x `Union` x

-- flatten 函数将 Line 替换为一个空格，把两行变成一行。
flatten :: Doc -> Doc
flatten (x `Concat` y) = flatten x `Concat` flatten y
flatten Line = Char ' '
-- 取出 Union 的 left
flatten (x `Union` _) = flatten x
flatten other = other

-- 当左括号（无论是小括号，中括号，还是大括号）出现时， 之后的行应该缩进，直到对应的右括号出现为止。
nest :: Int -> Doc -> Doc
-- 缩进为 n 个 字符
nest n xy@(x `Concat` y)
  | nest_isStt x = nest n x <> nest (2 * n) y
  | nest_isEnd y = nest n x <> (text (replicate n '\b') <> y)
  | otherwise = nest n x <> nest n y
  where
    nest_isStt x = x `elem` [Char '(', Char '[', Char '{']
    nest_isEnd x = x `elem` [Char ')', Char ']', Char '}']
    nest_min id = min 0 (id - 1)
nest n (x `Union` Line) = x `Union` (Line <> text (replicate n ' '))
nest n x = x

prettyNest :: Int -> Int -> Doc -> String
prettyNest x y = pretty x . nest y

-- 紧凑转换函数
compact :: Doc -> String
compact x = transform [x]
  where
    transform [] = ""
    transform (d : ds) =
      case d of
        Empty -> transform ds
        Char c -> c : transform ds
        Text s -> s ++ transform ds
        Line -> '\n' : transform ds
        a `Concat` b -> transform (a : b : ds)
        _ `Union` b -> transform (b : ds)

-- pretty 多了一个参数：每行的最大宽度(有几列)。(假设我们使用等宽字体。)
pretty :: Int -> Doc -> String
pretty width x = best 0 [x]
  where
    -- col 记录当前列数, width 是允许最大列数
    best col (d : ds) = case d of
      Empty -> best col ds
      Char c -> c : best (col + 1) ds
      Text s -> s ++ best (col + length s) ds
      Line -> '\n' : best 0 ds
      a `Concat` b -> best col (a : b : ds)
      a `Union` b ->
        nicest
          col
          (best col (a : ds))
          (best col (b : ds))
    best _ _ = ""
    nicest col a b
      | (width - least) `fits` a = a
      | otherwise = b
      where
        least = min width col

-- ---- int 是否能容下 String
fits :: Int -> String -> Bool
w `fits` _ | w < 0 = False
w `fits` "" = True
w `fits` ('\n' : _) = True
w `fits` (c : cs) = (w - 1) `fits` cs